#include "agxv_private.h"
#include "agxv_entrypoints.h"
#include "vk_alloc.h"
#include "agx_device.h"

static const struct vk_instance_extension_table instance_extensions = {
   .KHR_device_group_creation           = true,
#ifdef VK_USE_PLATFORM_DISPLAY_KHR
   .KHR_display                         = true,
#endif
   .KHR_external_fence_capabilities     = true,
   .KHR_external_memory_capabilities    = true,
   .KHR_external_semaphore_capabilities = true,
   .KHR_get_display_properties2         = true,
   .KHR_get_physical_device_properties2 = true,
#ifdef VK_USE_PLATFORM_WAYLAND_KHR
   .KHR_wayland_surface                 = true,
#endif
#ifdef VK_USE_PLATFORM_XCB_KHR
   .KHR_xcb_surface                     = true,
#endif
#ifdef VK_USE_PLATFORM_XLIB_KHR
   .KHR_xlib_surface                    = true,
#endif
   .EXT_debug_report                    = true,
};

VKAPI_ATTR VkResult VKAPI_CALL
agxv_EnumerateInstanceExtensionProperties(const char *pLayerName,
                                          uint32_t *pPropertyCount,
                                          VkExtensionProperties *pProperties)
{
    /* We don't support any layers  */
    if (pLayerName)
        return vk_error(NULL, VK_ERROR_LAYER_NOT_PRESENT);

    return vk_enumerate_instance_extension_properties(
        &instance_extensions, pPropertyCount, pProperties);
}

VKAPI_ATTR void VKAPI_CALL
agxv_DestroyInstance(VkInstance _instance,
                     const VkAllocationCallbacks *pAllocator)
{
   AGXV_FROM_HANDLE(agxv_instance, instance, _instance);

   if (!instance)
      return;

   if (instance->physicalDeviceCount > 0) {
      /* We support at most one physical device. */
      assert(instance->physicalDeviceCount == 1);
      vk_physical_device_finish(&instance->physicalDevice.vk);
   }

   vk_instance_finish(&instance->vk);
   vk_free(&instance->vk.alloc, instance);
}

VKAPI_ATTR VkResult VKAPI_CALL
agxv_CreateInstance(const VkInstanceCreateInfo *pCreateInfo,
                     const VkAllocationCallbacks *pAllocator,
                     VkInstance *pInstance)
{
    struct agxv_instance *instance;
    VkResult result;

    assert(pCreateInfo->sType == VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO);

    if (pAllocator == NULL)
        pAllocator = vk_default_allocator();

    instance = vk_alloc(pAllocator, sizeof(*instance), 8,
                        VK_SYSTEM_ALLOCATION_SCOPE_INSTANCE);
    if (!instance)
        return vk_error(NULL, VK_ERROR_OUT_OF_HOST_MEMORY);

    struct vk_instance_dispatch_table dispatch_table;
    vk_instance_dispatch_table_from_entrypoints(&dispatch_table, &agxv_instance_entrypoints, true);

    result = vk_instance_init(&instance->vk,
                              &instance_extensions,
                              &dispatch_table,
                              pCreateInfo, pAllocator);

    if (result != VK_SUCCESS) {
        vk_free(pAllocator, instance);
        return vk_error(instance, result);
    }

    instance->physicalDeviceCount = -1;

    *pInstance = agxv_instance_to_handle(instance);

    return VK_SUCCESS;
}

static VkResult
enumerate_devices(struct agxv_instance *instance)
{
    VkResult result;
    struct agxv_physical_device *physicalDevice = &instance->physicalDevice;

    struct vk_physical_device_dispatch_table dispatch_table;
    vk_physical_device_dispatch_table_from_entrypoints(&dispatch_table, &agxv_physical_device_entrypoints, true);

    result = vk_physical_device_init(&physicalDevice->vk, &instance->vk, NULL,
                                     &dispatch_table);

    if (result != VK_SUCCESS)
        return result;

    if(!agx_open_device(instance, &physicalDevice->dev)) {
        return VK_ERROR_DEVICE_LOST;
    }

    instance->physicalDeviceCount = 1;

    return VK_SUCCESS;
}

VKAPI_ATTR VkResult VKAPI_CALL
agxv_EnumeratePhysicalDevices(VkInstance _instance,
                               uint32_t *pPhysicalDeviceCount,
                               VkPhysicalDevice *pPhysicalDevices)
{
    AGXV_FROM_HANDLE(agxv_instance, instance, _instance);
    VK_OUTARRAY_MAKE(out, pPhysicalDevices, pPhysicalDeviceCount);

    if (instance->physicalDeviceCount < 0) {
        VkResult result = enumerate_devices(instance);
        if (result != VK_SUCCESS &&
            result != VK_ERROR_INCOMPATIBLE_DRIVER)
            return result;
    }

    if (instance->physicalDeviceCount == 0)
        return VK_SUCCESS;

    assert(instance->physicalDeviceCount == 1);
    vk_outarray_append(&out, i) {
        *i = agxv_physical_device_to_handle(&instance->physicalDevice);
    }

    return vk_outarray_status(&out);
}

VKAPI_ATTR VkResult VKAPI_CALL
agxv_CreateDevice(VkPhysicalDevice physicalDevice,
                   const VkDeviceCreateInfo *pCreateInfo,
                   const VkAllocationCallbacks *pAllocator,
                   VkDevice *pDevice)
{
    AGXV_FROM_HANDLE(agxv_physical_device, physical_device, physicalDevice);
    struct agxv_device *device;
    struct agxv_instance *instance = (struct agxv_instance *)physical_device->vk.instance;
    VkResult result;

    assert(pCreateInfo->sType == VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO);

    if (pCreateInfo->pEnabledFeatures) {
    }

    device = vk_zalloc2(&physical_device->vk.instance->alloc, pAllocator,
                        sizeof(*device), 8,
                        VK_SYSTEM_ALLOCATION_SCOPE_DEVICE);
    if (!device)
        return vk_error(instance, VK_ERROR_OUT_OF_HOST_MEMORY);

    struct vk_device_dispatch_table dispatch_table;
    vk_device_dispatch_table_from_entrypoints(&dispatch_table,
                                              &agxv_device_entrypoints, true);
    result = vk_device_init(&device->vk, &physical_device->vk,
                            &dispatch_table, pCreateInfo, pAllocator);
    if (result != VK_SUCCESS)
    {
        vk_free(&device->vk.alloc, device);
        return vk_error(instance, result);
    }

    device->instance = instance;
    device->pdevice = physical_device;

    if (pAllocator)
        device->vk.alloc = *pAllocator;
    else
        device->vk.alloc = physical_device->vk.instance->alloc;

    *pDevice = agxv_device_to_handle(device);

    return VK_SUCCESS;
}

VKAPI_ATTR VkResult VKAPI_CALL
agxv_CreateBuffer(VkDevice _device,
                   const VkBufferCreateInfo *pCreateInfo,
                   const VkAllocationCallbacks *pAllocator,
                   VkBuffer *pBuffer)
{
    AGXV_FROM_HANDLE(agxv_device, device, _device);
    struct agxv_buffer *buffer;

    assert(pCreateInfo->sType == VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO);

    buffer = vk_object_alloc(&device->vk, pAllocator, sizeof(*buffer),
                             VK_OBJECT_TYPE_BUFFER);
    if (buffer == NULL)
        return vk_error(device->instance, VK_ERROR_OUT_OF_HOST_MEMORY);

    buffer->size = pCreateInfo->size;
    buffer->usage = pCreateInfo->usage;
    buffer->flags = pCreateInfo->flags;

    *pBuffer = agxv_buffer_to_handle(buffer);

    return VK_SUCCESS;
}

VKAPI_ATTR VkResult VKAPI_CALL
agxv_CreateFramebuffer(VkDevice _device,
                       const VkFramebufferCreateInfo *pCreateInfo,
                       const VkAllocationCallbacks *pAllocator,
                       VkFramebuffer *pFramebuffer)
{
    AGXV_FROM_HANDLE(agxv_device, device, _device);
    struct agxv_framebuffer *framebuffer;

    assert(pCreateInfo->sType == VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO);

    size_t size = sizeof(*framebuffer) +
                  sizeof(struct agxv_image_view *) * pCreateInfo->attachmentCount;
    framebuffer = vk_object_zalloc(&device->vk, pAllocator, size,
                                   VK_OBJECT_TYPE_FRAMEBUFFER);
    if (framebuffer == NULL)
        return vk_error(device->instance, VK_ERROR_OUT_OF_HOST_MEMORY);

    framebuffer->width = pCreateInfo->width;
    framebuffer->height = pCreateInfo->height;
    framebuffer->layers = pCreateInfo->layers;

    framebuffer->attachment_count = pCreateInfo->attachmentCount;
    framebuffer->color_attachment_count = 0;
    for (uint32_t i = 0; i < pCreateInfo->attachmentCount; i++) {
        framebuffer->attachments[i] =
            agxv_image_view_from_handle(pCreateInfo->pAttachments[i]);
        if (framebuffer->attachments[i]->aspects & VK_IMAGE_ASPECT_COLOR_BIT)
            framebuffer->color_attachment_count++;
    }

    *pFramebuffer = agxv_framebuffer_to_handle(framebuffer);

    return VK_SUCCESS;
}

VKAPI_ATTR VkResult VKAPI_CALL
agxv_AllocateMemory(VkDevice _device,
                    const VkMemoryAllocateInfo *pAllocateInfo,
                    const VkAllocationCallbacks *pAllocator,
                    VkDeviceMemory *pMem)
{
    AGXV_FROM_HANDLE(agxv_device, device, _device);
    struct agxv_device_memory *mem;
    struct agx_bo *bo;

    assert(pAllocateInfo->sType == VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO);

    if (pAllocateInfo->allocationSize == 0) {
        /* Apparently, this is allowed */
        *pMem = VK_NULL_HANDLE;
        return VK_SUCCESS;
    }

    mem = vk_object_alloc(&device->vk, pAllocator, sizeof(*mem),
                          VK_OBJECT_TYPE_DEVICE_MEMORY);
    if (mem == NULL)
        return vk_error(device->instance, VK_ERROR_OUT_OF_HOST_MEMORY);

    const VkImportMemoryFdInfoKHR *fd_info =
        vk_find_struct_const(pAllocateInfo->pNext, IMPORT_MEMORY_FD_INFO_KHR);
    if (fd_info && !fd_info->handleType)
        fd_info = NULL;

    if (fd_info) {
        assert(fd_info->handleType == VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT ||
               fd_info->handleType == VK_EXTERNAL_MEMORY_HANDLE_TYPE_DMA_BUF_BIT_EXT
               );

        assert(0);
    } else {
        bo = agx_bo_create(&device->pdevice->dev, pAllocateInfo->allocationSize, AGX_MEMORY_TYPE_FRAMEBUFFER);
        if (!bo) {
            vk_object_free(&device->vk, pAllocator, mem);
            return VK_ERROR_OUT_OF_DEVICE_MEMORY;
        }

        mem->bo = bo;
    }

    *pMem = agxv_device_memory_to_handle(mem);

    return VK_SUCCESS;
}

VKAPI_ATTR void VKAPI_CALL
agxv_FreeMemory(VkDevice _device,
              VkDeviceMemory _mem,
              const VkAllocationCallbacks *pAllocator)
{
   AGXV_FROM_HANDLE(agxv_device, device, _device);
   AGXV_FROM_HANDLE(agxv_device_memory, mem, _mem);

   if (mem == NULL)
      return;

   agx_bo_unreference(mem->bo);
   vk_object_free(&device->vk, pAllocator, mem);
}

VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL
agxv_GetInstanceProcAddr(VkInstance _instance,
                         const char *pName)
{
    AGXV_FROM_HANDLE(agxv_instance, instance, _instance);
    return vk_instance_get_proc_addr(&instance->vk,
                                     &agxv_instance_entrypoints,
                                     pName);
}

/* With version 1+ of the loader interface the ICD should expose
 * vk_icdGetInstanceProcAddr to work around certain LD_PRELOAD issues seen in apps.
 */
PUBLIC
VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL
vk_icdGetInstanceProcAddr(VkInstance instance,
                          const char *pName);

PUBLIC
VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL
vk_icdGetInstanceProcAddr(VkInstance instance,
                          const char *pName)
{
    return agxv_GetInstanceProcAddr(instance, pName);
}

/* With version 4+ of the loader interface the ICD should expose
 * vk_icdGetPhysicalDeviceProcAddr()
 */
PUBLIC
VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL
vk_icdGetPhysicalDeviceProcAddr(VkInstance _instance,
                                const char *pName);

PFN_vkVoidFunction
vk_icdGetPhysicalDeviceProcAddr(VkInstance _instance,
                                const char *pName)
{
    AGXV_FROM_HANDLE(agxv_instance, instance, _instance);
    return vk_instance_get_physical_device_proc_addr(&instance->vk, pName);
}
